import os
import gdown
import zipfile
import argparse
from typing import List, Optional

def parse_download_links(links_file: Optional[str] = None, gdrive_link: Optional[str] = None) -> List[str]:
    if links_file is None and gdrive_link is None:
        raise ValueError("Either links_file or gdrive_link must be provided")
    if gdrive_link:
        return [gdrive_link]
    with open(links_file, 'r') as f:
        links = f.readlines()
    urls: List[str] = []
    for line in links:
        line = line.strip()
        if not line:
            continue
        if ': ' in line:
            _, url = line.split(': ', 1)
        else:
            url = line
        urls.append(url)
    return urls

def download_and_extract_data(output_path: str, urls: List[str], unzip: bool = True, remove_zip: bool = True) -> None:
    os.makedirs(output_path, exist_ok=True)
    if not output_path.endswith('/'):
        output_path += '/'
    for url in urls:
        print(f'>>>>> Downloading {url} to {output_path}')
        output = gdown.download(url=url, output=output_path, fuzzy=True, quiet=False)
        print(f'Downloaded to {output}')
        if unzip and zipfile.is_zipfile(output):
            print('Unzipping file...')
            with zipfile.ZipFile(output, 'r') as zip_ref:
                zip_ref.extractall(output_path)
            print(f'Unzipped to {output_path}')
        if remove_zip and os.path.exists(output) and zipfile.is_zipfile(output):
            os.remove(output)
            print(f'Removed zip file {output}')

def main():
    parser = argparse.ArgumentParser(description='Download and extract aerial-megadepth data')
    parser.add_argument('--output_path', type=str, required=True,
                        help='Path where the extracted data should be saved')
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument('--gdrive_links_file', type=str,
                        help='Path to text file containing filename and Google Drive links')
    group.add_argument('--gdrive_link', type=str,
                        help='Single Google Drive link to download')

    parser.add_argument('--unzip', dest='unzip', action='store_true', help='Unzip the downloaded files')
    parser.add_argument('--no-unzip', dest='unzip', action='store_false', help='Do not unzip')
    parser.set_defaults(unzip=True)

    parser.add_argument('--remove_zip', dest='remove_zip', action='store_true', help='Remove the zip files after extraction')
    parser.add_argument('--keep-zip', dest='remove_zip', action='store_false', help='Keep the zip files after extraction')
    parser.set_defaults(remove_zip=True)

    args = parser.parse_args()

    urls = parse_download_links(args.gdrive_links_file, args.gdrive_link)
    download_and_extract_data(args.output_path, urls, args.unzip, args.remove_zip)

if __name__ == '__main__':
    main()
